from deepsignal3.utils import bam_reader
import pod5
import os
from collections import defaultdict
from uuid import UUID

def extract_overlapping_reads_from_folder(bam_path, pod5_folder, output_pod5_path, max_reads=5):
    # 创建 ReadIndexedBam 对象
    bam_file = bam_reader.ReadIndexedBam(bam_path)
    bam_file.open()

    # 创建一个集合来存储有重叠的 read_id
    overlapping_reads = set()

    # 使用字典存储参考序列上的覆盖区间，key 是参考名称，value 是一个列表（覆盖区间）
    reference_coverage = defaultdict(list)

    # 找到的满足条件的重叠区域的 reads 集合
    final_overlapping_reads = set()

    # 遍历 BAM 文件，找到重叠的 reads
    for read in bam_file:
        if read.is_unmapped:
            continue

        ref_name = read.reference_name
        start = read.reference_start
        end = read.reference_end
        read_id = read.query_name

        # 遍历当前参考基因组上的所有区域，检查与当前 read 是否有重叠
        overlapping_regions = []
        for region in reference_coverage[ref_name]:
            region_start, region_end, region_reads = region

            # 如果有重叠，更新重叠区信息
            if start < region_end and end > region_start:
                # 更新重叠区范围
                region_start = min(region_start, start)
                region_end = max(region_end, end)
                region_reads.add(read_id)
                overlapping_regions.append((region_start, region_end, region_reads))
                final_overlapping_reads = region_reads

        # 更新参考覆盖区，将更新后的区域重新赋值
        reference_coverage[ref_name] = [r for r in reference_coverage[ref_name] if r not in overlapping_regions]
        reference_coverage[ref_name].extend(overlapping_regions)

        # 如果没有与任何现有区域重叠的情况，创建一个新的覆盖区域
        if not overlapping_regions:
            reference_coverage[ref_name].append((start, end, {read_id}))

        # 检查是否有重叠区域包含至少 5 条 reads，如果是，停止遍历
        for region in reference_coverage[ref_name]:
            if len(region[2]) >= 10:
                final_overlapping_reads = region[2]
                print(f"Read {read_id} overlaps with at least 5 reads in reference {ref_name}.")
                break
        if len(final_overlapping_reads) >= 10:
            print(f"Found at least 10 overlapping reads in {bam_path}.")
            break

    # 遍历 pod5 文件夹中的所有 pod5 文件，提取相应的 reads
    for read_id in final_overlapping_reads:
        print(f"Extracting read {read_id} from pod5 files...")
    if final_overlapping_reads:
        print(f"Found overlapping region with reads: {final_overlapping_reads}")

        # 遍历 pod5 文件夹中的所有 pod5 文件，提取相应的 reads
        with pod5.Writer(output_pod5_path) as writer:
            reads_to_write = final_overlapping_reads.copy()  # 保留原始集合，避免过早清空
            for pod5_file_name in os.listdir(pod5_folder):
                pod5_file_path = os.path.join(pod5_folder, pod5_file_name)

                if os.path.isfile(pod5_file_path) and pod5_file_name.endswith('.pod5'):
                    with pod5.Reader(pod5_file_path) as reader:
                        for read_record in reader:
                            read_id = str(read_record.read_id)  # 假设 ReadRecord 有 read_id 属性
                            if read_id in reads_to_write:
                                read = pod5.Read(
                                    read_id=read_record.read_id,  # 直接使用原始 UUID 对象
                                    read_number=read_record.read_number,
                                    start_sample=read_record.start_sample,
                                    median_before=read_record.median_before,
                                    end_reason=read_record.end_reason,
                                    calibration=read_record.calibration,
                                    pore=read_record.pore,
                                    run_info=read_record.run_info,
                                    signal=read_record.signal,
                                )
                                writer.add_read(read)  # 将 ReadRecord 直接写入新的文件
                                reads_to_write.remove(read_id)  # Remove after writing to avoid redundant reads
                                print(f"Written read: {read_id}")
                                if len(reads_to_write) <= 0:
                                    break

                # 如果已经找到所有需要的 read，就停止遍历 pod5 文件
                if len(reads_to_write) <= 0:
                    break

        print(f"Successfully written selected overlapping reads to {output_pod5_path}")

    else:
        print("No overlapping region with at least 10 reads found.")

# 使用示例
bam_path = "/home/xiaoyifu/data/HG003/hg003.demo.bam"
pod5_folder = "/home/xiaoyifu/data/HG003/demo_pod5/"
output_pod5_path = "/home/xiaoyifu/data/HG003/demo_output.pod5"
extract_overlapping_reads_from_folder(bam_path, pod5_folder, output_pod5_path)
